package com.atbeijing.D09;

import org.apache.flink.api.common.eventtime.SerializableTimestampAssigner;
import org.apache.flink.api.common.eventtime.WatermarkStrategy;
import org.apache.flink.api.common.functions.MapFunction;
import org.apache.flink.api.common.state.ValueState;
import org.apache.flink.api.common.state.ValueStateDescriptor;
import org.apache.flink.api.common.typeinfo.Types;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.api.java.tuple.Tuple3;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.KeyedProcessFunction;
import org.apache.flink.util.Collector;

import java.util.Random;

/**
 * 每10秒分组求一次和
 * 数据倾斜问题
 */
public class Example7 {
    public static void main(String[] args) throws Exception {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        env.setParallelism(10);

        env
                .fromElements(
                        //(key,次数,时间戳)
                        Tuple3.of("a", 1L, 1000L),
                        Tuple3.of("a", 1L, 2000L),
                        Tuple3.of("a", 1L, 3000L),
                        Tuple3.of("a", 1L, 4000L),
                        Tuple3.of("a", 1L, 5000L),
                        Tuple3.of("a", 1L, 6000L),
                        Tuple3.of("a", 1L, 7000L),
                        Tuple3.of("a", 1L, 8000L),
                        Tuple3.of("a", 1L, 9000L),
                        Tuple3.of("a", 1L, 10000L),
                        Tuple3.of("b", 1L, 11000L)
                )
                //转换数据结构: key+随机数,避免keyBy数据倾斜
                .map(new MapFunction<Tuple3<String, Long, Long>, Tuple3<String, Long, Long>>() {
                    @Override
                    public Tuple3<String, Long, Long> map(Tuple3<String, Long, Long> value) throws Exception {
                        Random random = new Random();
                        //(a~0,1,2000)  (a~2,1,3000)
                        return Tuple3.of(value.f0+"~"+random.nextInt(4),value.f1,value.f2);
                    }
                })
                //添加水位线
                .assignTimestampsAndWatermarks(
                        WatermarkStrategy.<Tuple3<String, Long, Long>>forMonotonousTimestamps()
                        .withTimestampAssigner(new SerializableTimestampAssigner<Tuple3<String, Long, Long>>() {
                            @Override
                            public long extractTimestamp(Tuple3<String, Long, Long> element, long recordTimestamp) {
                                return element.f2;
                            }
                        })
                )
                //分流
                .keyBy(r -> r.f0)
                //处理每条流
                .process(new KeyedProcessFunction<String, Tuple3<String, Long, Long>, Tuple2<String,Long>>() {
                    //存放数据流的处理结果(key,sum)
                    private ValueState<Tuple2<String,Long>> sum;
                    //存放定时器时间戳
                    private ValueState<Long> timers;

                    @Override
                    public void open(Configuration parameters) throws Exception {
                        super.open(parameters);
                        //初始化状态变量
                        sum=getRuntimeContext().getState(new ValueStateDescriptor<Tuple2<String,Long>>("sum",Types.TUPLE(Types.STRING,Types.LONG)));
                        timers=getRuntimeContext().getState(new ValueStateDescriptor<Long>("timers",Types.LONG));
                    }

                    //来一个数据处理一次 (a~0,1,2000)  (a~2,1,3000)
                    @Override
                    public void processElement(Tuple3<String, Long, Long> value, Context ctx, Collector<Tuple2<String, Long>> out) throws Exception {
                        //每10秒输出一次,所以要有一个10的定时器
                        //第一条数据
                        if (sum.value()==null ){
                            sum.update(Tuple2.of(value.f0,value.f1));
                            //10秒后的定时器
                            ctx.timerService().registerEventTimeTimer(value.f2+10*1000L);
                            timers.update(value.f2+10*1000L);
                        }else {
                            //累加
                            Long acc = sum.value().f1;
                            sum.update(Tuple2.of(value.f0,acc+value.f1));
                            //判断定时器有没有触发,如果触发了,就再设置一个10秒后的定时器
                            if (timers.value()==null){
                                ctx.timerService().registerEventTimeTimer(value.f2+10*1000L);
                                timers.update(value.f2+10*1000L);
                            }
                        }
                    }

                    //定时器触发后输出数据,清空定时器
                    @Override
                    public void onTimer(long timestamp, OnTimerContext ctx, Collector<Tuple2<String, Long>> out) throws Exception {
                        super.onTimer(timestamp, ctx, out);
                        Tuple2<String, Long> value = sum.value();
                        out.collect(value);
                        timers.clear();
                    }
                })//(a~3,3) (a~0,2) (a~2,2) (b~2,1) (a~1,3)
                //将key还原再聚合求出真实结果
                .map(new MapFunction<Tuple2<String, Long>, Tuple2<String, Long>>() {
                    @Override
                    public Tuple2<String, Long> map(Tuple2<String, Long> value) throws Exception {
                        String s = value.f0.split("~")[0];
                        return Tuple2.of(s,value.f1);
                    }
                })//(a,3) (a,2) (a,2) (b,1) (a,3)
                .keyBy(r -> r.f0)
                //聚合
                .process(new KeyedProcessFunction<String, Tuple2<String, Long>, Tuple2<String, Long>>() {
                    //存放数据流的处理结果(key,sum)
                    private ValueState<Tuple2<String,Long>> sum;

                    @Override
                    public void open(Configuration parameters) throws Exception {
                        super.open(parameters);
                        sum=getRuntimeContext().getState(new ValueStateDescriptor<Tuple2<String, Long>>("sum",Types.TUPLE(Types.STRING,Types.LONG)));
                    }

                    @Override
                    public void processElement(Tuple2<String, Long> value, Context ctx, Collector<Tuple2<String, Long>> out) throws Exception {
                        if (sum.value()==null){
                            sum.update(Tuple2.of(value.f0,value.f1));
                        }else {
                            sum.update(Tuple2.of(value.f0,value.f1+sum.value().f1));
                        }
                        Tuple2<String, Long> sumResult = sum.value();
                        out.collect(sumResult);
                    }
                })
                .print();

        env.execute();
    }
}
